#!/usr/bin/env python3
"""
RAG项目启动脚本
功能：自动检查并构建索引（如果需要），然后启动FastAPI服务
"""

import sys
import subprocess
import json
from pathlib import Path

def check_requirements() -> bool:
    """检查依赖是否安装"""
    try:
        print("✓ 依赖检查通过")
        return True
    except ImportError as e:
        print(f"✗ 缺少依赖: {e}")
        print("请运行: pip install -r requirements.txt")
        return False

def check_index_exists() -> bool:
    """检查索引文件是否存在"""
    index_dir = Path("indexes")
    if not index_dir.exists():
        print("✗ 索引目录不存在")
        return False
    
    required_files = ["vectors.faiss", "metadatas.pkl", "config.json"]
    missing_files = []
    
    for file in required_files:
        file_path = index_dir / file
        if not file_path.exists():
            missing_files.append(file)
    
    if missing_files:
        print(f"✗ 缺少索引文件: {', '.join(missing_files)}")
        return False
    
    print("✓ 索引文件存在")
    return True

def check_model_exists() -> bool:
    """检查模型是否已下载"""
    try:
        # 读取配置文件
        config_path = Path("indexes/config.json")
        if not config_path.exists():
            print("✗ 配置文件不存在")
            return False
        
        with open(config_path, "r", encoding="utf-8") as f:
            cfg = json.load(f)


        return True
    except Exception as e:
        print(f"✗ 无法检查模型文件: {e}")
        return False

def build_index() -> bool:
    """构建索引"""
    try:
        print("开始构建索引...")
        result = subprocess.run([sys.executable, "build_index.py"], 
                              capture_output=True, text=True, encoding='utf-8')
        
        if result.returncode == 0:
            print("✓ 索引构建成功")
            return True
        else:
            print(f"✗ 索引构建失败: {result.stderr}")
            return False
    except Exception as e:
        print(f"✗ 构建索引时出错: {e}")
        return False

def start_server() -> bool:
    """启动服务"""
    try:
        print("启动FastAPI服务...")
        print("服务将在 http://localhost:8000 启动")
        print("API文档地址: http://localhost:8000/docs")
        print("按 Ctrl+C 停止服务")
        
        # 使用uvicorn启动服务
        subprocess.run([
            "uvicorn", 
            "server:app", 
            "--reload", 
            "--host", "0.0.0.0", 
            "--port", "8000"
        ])
        
        return True
    except KeyboardInterrupt:
        print("\n服务已停止")
        return True
    except Exception as e:
        print(f"✗ 启动服务失败: {e}")
        return False

def main():
    """主函数"""
    print("=== RAG项目启动器 ===")
    
    # 检查依赖
    if not check_requirements():
        return 1
    
    # 检查索引
    if not check_index_exists():
        print("检测到缺少索引文件，开始构建...")
        if not build_index():
            return 1
    else:
        # 检查模型
        if not check_model_exists():
            print("检测到模型文件缺失，开始构建索引...")
            if not build_index():
                return 1
    
    # 启动服务
    if not start_server():
        return 1
    
    return 0

if __name__ == "__main__":
    sys.exit(main())